// Copyright (c) 2013, Hannes Würfel <hannes.wuerfel@student.hpi.uni-potsdam.de>
// Computer Graphics Systems Group at the Hasso-Plattner-Institute, Germany
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//   * Redistributions of source code must retain the above copyright notice,
//     this list of conditions and the following disclaimer.
//   * Redistributions in binary form must reproduce the above copyright
//     notice, this list of conditions and the following disclaimer in the
//     documentation and/or other materials provided with the distribution.
//   * Neither the name of the Computer Graphics Systems Group at the
//     Hasso-Plattner-Institute (HPI), Germany nor the names of its
//     contributors may be used to endorse or promote products derived from
//     this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

// the Per-Pixel Linked Lists construction was taken from the book GPU Pro2 

#include "Lighting.fx"
#include "LinkedList.fx"

struct VertexIn
{
	float3 PosL    : POSITION;
	float3 NormalL : NORMAL;
    float2 Tex     : TEXCOORD;
};

struct VertexOut
{
	float4 PosH    : SV_POSITION;
    float3 PosW    : POSITION;
    float3 NormalW : NORMAL;
    float2 Tex     : TEXCOORD;
};

Texture2D gTexture;

SamplerState anisotropicSampler
{
    Filter = ANISOTROPIC;
    MaxAnisotropy = 4;

    AddressU = WRAP;
    AddressV = WRAP;
};

cbuffer cbPerFrame
{
	DirectionalLight gDirLights[3];
	float3 gEyePosW;
};

cbuffer cbPerObject
{
	float4x4 gWorld;
	float4x4 gWorldInvTranspose;
	float4x4 gWorldViewProj;
    float4x4 gTexTransform;
	Material gMaterial;
};

cbuffer cbPerViewportUpdate
{
	int gViewportWidth;
};

VertexOut vertexShader(VertexIn vin)
{
	VertexOut vout;
	
	vout.PosW    = mul(float4(vin.PosL, 1.0f), gWorld).xyz;
	vout.NormalW = mul(vin.NormalL, (float3x3)gWorldInvTranspose);
		
	vout.PosH = mul(float4(vin.PosL, 1.0f), gWorldViewProj);
    vout.Tex  = mul(float4(vin.Tex, 0.0f, 1.0f), gTexTransform).xy;
	
	return vout;
}

// uav declaration
RWByteAddressBuffer headpointerBuf    : register(u1);
RWStructuredBuffer<NodeData> nodesBuf : register(u2);

// depth and stencil test before fragment shader
[earlydepthstencil]
float4 storeFragments(VertexOut pin,
                     uniform bool gUseLighting,
                     uniform bool gUseTexture) : SV_Target
{
	// phong shading
	pin.NormalW = normalize(pin.NormalW);
	float3 toEye = gEyePosW - pin.PosW;

	float distToEye = length(toEye); 
	toEye /= distToEye;

	float4 ambient = float4(0.0f, 0.0f, 0.0f, 0.0f);
	float4 diffuse = float4(0.0f, 0.0f, 0.0f, 0.0f);
	float4 spec    = float4(0.0f, 0.0f, 0.0f, 0.0f);

    if(gUseLighting)
    {
	    float4 A, D, S;
	    computeDirectionalLight(gMaterial,
							    gDirLights[0],
							    pin.NormalW,
							    toEye, 
							    A,
							    D,
							    S);

	    ambient += A;
	    diffuse += D;
	    spec    += S;
    }

    float4 litColor = ambient + diffuse + spec;

    if(gUseTexture)
    {
        litColor += gTexture.Sample(anisotropicSampler, pin.Tex);
    }

	litColor.a = gMaterial.diffuse.a;
	
	//  writing to per pixel linked list
	NodeData node;
	node.color = packFloat4IntoUint(litColor);
	node.depth = packDepthIntoUint(pin.PosH.z);
	
	// retrieve current pixel count and increase counter
	uint pixelCount = nodesBuf.IncrementCounter();
	
	// convert fragment 2D coordinates to byte linear address
	uint2 screenPos = uint2(pin.PosH.xy);
	uint linearAddressInBytes = 4 * (screenPos.y * gViewportWidth + screenPos.x);
	
	// pixelCount becomes the new entry point of the linked list for this pixel
	// the previous head pointer becomes the next link  of the current node
	headpointerBuf.InterlockedExchange(linearAddressInBytes, pixelCount, node.next);
	
	// add new fragment in nodeBuf
	nodesBuf[pixelCount] = node;
	
	// because no render target was bound, this has no effect
	return float4(0.0f, 0.0f, 0.0f, 0.0f);
}

// shader resource views
Buffer<uint> headPointerBufSRV 			: register(t0);
StructuredBuffer<NodeData> nodesBufSRV	: register(t1);
Texture2D gBackgroundTex				: register(t3); // ?

// fragment sorting array (for insertion sort, TODO try to implement bitonic sort
#define MAX_SORTED_FRAGMENTS 18
static uint2 gSortedFragments[MAX_SORTED_FRAGMENTS+1];

float4 sortFragments(VertexOut pin) : SV_Target
{
	// convert pixel 2D coordinates to linear address
	uint2 screenPos = uint2(pin.PosH.xy);
	uint linearAddress = (screenPos.y * gViewportWidth + screenPos.x);
	
	// fetch offset of the first fragment for the current pixel
	uint offset = headPointerBufSRV[linearAddress];
	
	int nFragments = 0;
	while(offset != 0xffffffff)
	{
		// retrieve fragment
		NodeData node = nodesBufSRV[offset];
		
		// copy data into sorting array
		gSortedFragments[nFragments] = uint2(node.color, node.depth);
		
		// insertion sort (depth values front to back)
		int i = nFragments;
		while((i > 0) &&
			  (gSortedFragments[max(i - 1, 0)].y > 
			   gSortedFragments[i].y))
		{
			uint2 tmp = gSortedFragments[i];
			gSortedFragments[i] = gSortedFragments[max(i - 1, 0)];
			gSortedFragments[max(i - 1, 0)] = tmp;
			--i;
		}
		
		nFragments = min(nFragments + 1, MAX_SORTED_FRAGMENTS);
		offset = node.next;
	}
	
	// retrieve current color from background (opaqua and alpha tested scene rendered to RTV before)
	float4 outColor = gBackgroundTex.Sample(anisotropicSampler, pin.Tex);
	
	// manual SRCALPHA-INVSRCALPHA blending
	
	for(int f = nFragments - 1; f >= 0; --f)
	{
		float4 color = unpackUintIntoFloat4(gSortedFragments[f].x);
		outColor.rgb = lerp(outColor.rgb, color.rgb, color.a);
	}
	
	return outColor;
}

technique11 StoreFragments
{
    pass P0
    {
        SetVertexShader(CompileShader(vs_5_0, vertexShader()));
        SetGeometryShader(NULL);
        SetPixelShader(CompileShader(ps_5_0, storeFragments(true, false)));
    }
}

technique11 SortFragments
{
	pass P0
	{
		SetVertexShader(CompileShader(vs_5_0, vertexShader()));
        SetGeometryShader(NULL);
        SetPixelShader(CompileShader(ps_5_0, sortFragments()));
	}
}